import os

from ruamel.yaml import YAML
from ruamel.yaml.comments import CommentedMap

import os
import shutil
import random
from glob import glob


def move_text_to_yolo_dir(source_dir, target_dir, txt_name, isCopy=False):
    txt_file = txt_name + '.txt'
    txt_path = os.path.join(source_dir, txt_file)
    if os.path.exists(txt_path):
        if isCopy:
            shutil.copy2(txt_path, os.path.join(target_dir, txt_file))
        else:
            shutil.move(txt_path, os.path.join(target_dir, txt_file))


def split_images_to_train_val(source_dir, train_ratio=0.8):
    # 定义图片文件扩展名（可根据需要补充）
    image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.gif', '*.tiff']
    # 定义目标目录
    train_dir = os.path.join(source_dir, 'data', 'images', 'train')
    lb_train_dir = os.path.join(source_dir, 'data', 'labels', 'train')

    val_dir = os.path.join(source_dir, 'data', 'images', 'val')
    lb_val_dir = os.path.join(source_dir, 'data', 'labels', 'val')

    # 收集源目录下所有图片文件的路径
    image_paths = []
    for ext in image_extensions:
        # 递归搜索所有子目录中的图片
        image_paths.extend(glob(os.path.join(source_dir, '*', ext), recursive=True))

    if not image_paths:
        print(f"错误：在 {source_dir} 中未找到任何图片文件")
        return False

    if len(image_paths) == 0:
        print('不存在图片文件')

    if len(image_paths) < 3:
        for img_path in image_paths:
            img_name = os.path.basename(img_path)
            train_dest_path = os.path.join(train_dir, img_name)
            val_dest_path = os.path.join(val_dir, img_name)

            name, ext = os.path.splitext(img_name)
            shutil.copy2(img_path, train_dest_path)  # 保留文件元数据
            shutil.move(img_path, val_dest_path)

            move_text_to_yolo_dir(source_dir, lb_train_dir, name, True)
            move_text_to_yolo_dir(source_dir, lb_val_dir, name)

    else:
        # 随机打乱图片顺序
        random.shuffle(image_paths)
        # 计算训练集和验证集的分割点
        split_idx = int(len(image_paths) * train_ratio)
        train_images = image_paths[:split_idx]
        val_images = image_paths[split_idx:]

        # 复制图片到训练集目录
        for img_path in train_images:
            img_name = os.path.basename(img_path)
            dest_path = os.path.join(train_dir, img_name)
            # 避免文件名冲突（如存在同名文件则添加后缀）
            counter = 1
            name, ext = os.path.splitext(img_name)
            while os.path.exists(dest_path):
                dest_path = os.path.join(train_dir, f"{name}_{counter}{ext}")
                counter += 1
            shutil.move(img_path, dest_path)
            move_text_to_yolo_dir(source_dir, lb_train_dir, name)

        # 复制图片到验证集目录
        for img_path in val_images:
            img_name = os.path.basename(img_path)
            dest_path = os.path.join(val_dir, img_name)
            counter = 1
            name, ext = os.path.splitext(img_name)
            while os.path.exists(dest_path):
                dest_path = os.path.join(val_dir, f"{name}_{counter}{ext}")
                counter += 1
            shutil.move(img_path, dest_path)
            move_text_to_yolo_dir(source_dir, lb_val_dir, name)

    print(f"分割完成")
    return True


def create_images(dir):
    # 定义要创建的目录列表
    directories = [
        "data/images/train",
        "data/images/val",
        "data/labels/train",
        "data/labels/val"
    ]

    # 遍历目录列表并创建
    for dir_path in directories:
        file_path = os.path.join(dir, dir_path)
        if os.path.exists(file_path) is False:
            # exist_ok=True 表示如果目录已存在，不抛出异常
            os.makedirs(file_path, exist_ok=True)
            print(f"已创建目录: {dir_path}")


def to_write_yolo_yaml(default_save_dir, label_hist):

    create_images(default_save_dir)
    split_images_to_train_val(default_save_dir)

    filename = os.path.join(default_save_dir, 'data.yaml')

    data = CommentedMap()
    data['path'] = default_save_dir
    data['train'] = default_save_dir
    data['val'] = default_save_dir

    data['nc'] = 4

    data['names'] = dict(enumerate(label_hist))

    # 添加注释
    data.yaml_set_comment_before_after_key('path', before='路径')
    data.yaml_set_comment_before_after_key('train', before='train路径')
    data.yaml_set_comment_before_after_key('val', before='val路径')

    data.yaml_set_comment_before_after_key('nc', before='物品个数')
    data.yaml_set_comment_before_after_key('names', before='物品名称')

    # 写入YAML文件
    yaml = YAML()
    yaml.encoding = 'utf-8'
    with open(filename, 'w', encoding='utf-8') as file:
        yaml.dump(data, file)
